import paramiko

from configs.settings import flask_config


class SSH:
    def __init__(self, timeout=30):
        storage = flask_config.STORAGE
        self.ip = storage["HOSTNAME"]
        self.port = storage["PORT"]
        self.username = storage["USERNAME"]
        self.password = storage["PW"]
        self.timeout = timeout
        self.ssh = paramiko.SSHClient()
        self.transport = paramiko.Transport(self.ip, self.port)

    def _password_connect(self):
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.ssh.connect(
            hostname=self.ip,
            port=self.port,
            username=self.username,
            password=self.password,
        )
        self.transport.connect(username=self.username, password=self.password)

    def connect(self):
        try:
            self._password_connect()
        except Exception as e:
            print("ssh password connect failed:", e)

    def close(self):
        self.transport.close()
        self.ssh.close()

    def execute_cmd(self, cmd):
        stdin, stdout, stderr = self.ssh.exec_command(cmd)
        res, err = stdout.read(), stderr.read()
        result = res if res else err
        return result.decode()

    def _sftp_get(self, remote_file, local_file):
        sftp = paramiko.SFTPClient.from_transport(self.transport)
        sftp.get(remote_file, local_file)

    def _sftp_put(self, local_file, remote_file):
        sftp = paramiko.SFTPClient.from_transport(self.transport)
        sftp.put(local_file, remote_file)

    def _sftp_put_stream(self, file, remote_file):
        sftp = paramiko.SFTPClient.from_transport(self.transport)
        sftp.putfo(file.stream, remote_file, confirm=True)


if __name__ == "__main__":
    ssh = SSH()
    ssh.connect()
    cmd = "ls -lh"
    ssh.execute_cmd(cmd)  # 执行命令

    remote_file, local_file = "xxx", "xxx"
    ssh._sftp_get(remote_file, local_file)  # 下载文件

    remote_dir, local_dir = "xxx", "xxx"
    ssh._sftp_put(remote_dir, remote_dir)  # 上传文件夹
    ssh.close()
